from dwave.system.composites import EmbeddingComposite
# from dwave.system.samplers import DWaveSampler
import neal
import dwavebinarycsp
from penalty_function import makeExactOne
from collections import defaultdict
# 资源
# 节点
# 时间

# 数据对应的节点
DATA_TABLE = {
    0: {
        'size': 1,
        'nodes': [],  # 哪都没有存
    },
    1: {
        'size': 1,
        'nodes': [1,2],
    },
    2: {
        'size': 1,
        'nodes': [1,5],
    },
    3: {
        'size': 1,
        'nodes': [1,6],
    },
}
DATA_NUMBER = len(DATA_TABLE)

TIME_FRAME_NUMBER = 100  # 假设 0.1 us 一段
TIME_FRAMES = [i for i in range(TIME_FRAME_NUMBER)]

NODE_NUMBER = 5
NODES = [i for i in range(NODE_NUMBER)]

WAIT_TIME_PAR_NODE = {
    node: node % 2
    for node in NODES    
}

MAX_SPEED = 10
NETWORK_SPEED = {
    (node1, node2): (node1 + node2) % 3 + 1
    for node1 in NODES
    for node2 in NODES    
}
for node in NODES:
    NETWORK_SPEED[(node, node)] = 10

# node读取时间、remote node节点的时间
def getNodeQubit(own_node, read_node, time_frame, data_resource):
    return f'N{own_node}_{read_node}_{time_frame}_{data_resource}'

sampler = neal.SimulatedAnnealingSampler()

# 每个节点的queue的数量 * 平均时间 要小
# 每个节点的利用率高（1的数量多）

# 每次都的时候要考虑 （传输速度，等待时间）,等待时常是考虑了之前的read的
# 每个节点应该发的节点是exact one的（未来可以试下有多个node，自己算一个传输快的）
Q = defaultdict(int)

exact_one_lagrange = 0.5
for time in TIME_FRAMES:
    for own_node in NODES:  #实际持有qubit的node
        for data in DATA_TABLE:
            qubits = [getNodeQubit(own_node, read_node, time, data) for read_node in NODES]
            makeExactOne(qubits, exact_one_lagrange)

# 假设当前定了，状态会变成这个定的


csp = dwavebinarycsp.ConstraintSatisfactionProblem(dwavebinarycsp.BINARY)

bqm = dwavebinarycsp.stitch(csp)

# https://link.springer.com/article/10.1007/s10951-019-00629-3
# online schedule algorithm